#!/usr/bin/env python3
# example: sudo env "PATH=$PATH" python cutter.py huobi okex
import os
import sys
import threading

urls = {
    'huobi': 'api.huobi.pro',
    'huobi_futures': 'api.hbdm.com',
    'okex': 'real.okex.com',
    'bitmex': 'www.bitmex.com'
}


def main():
  exchanges = sys.argv[1:]
  for exchange in exchanges:
    if exchange not in urls:
      raise RuntimeError(f"exchange [{exchange}] does not exist")

  for (exchange, url) in urls.items():
    if exchange in exchanges:
      print('cutting   ' + exchange)
      os.system(f'/usr/sbin/iptables -A INPUT -s {url} -j DROP >/dev/null 2>&1')
    else:
      print('restoring ' + exchange)
      os.system(f'/usr/sbin/iptables -D INPUT -s {url} -j DROP >/dev/null 2>&1')

  print('=====')
  print('done.')

  if len(exchanges) == 0:
    return

  # indefinite sleep
  try:
    print('press Ctrl+c to stop.')
    print('=====')
    forever = threading.Event()
    forever.wait()
  except Exception:
    pass
  finally:
    print()
    for (exchange, url) in urls.items():
      print('restoring   ' + exchange)
      os.system(f'/usr/sbin/iptables -D INPUT -s {url} -j DROP >/dev/null 2>&1')

    print('=====')
    print('restored')


if __name__ == "__main__":
  main()
